import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

class ReturnAnalysis:
    def __init__(self):
        self.codes=['600036','002594','300006','000905','002607']
    def get_return(self):
        #analysis array
        res_analysis=[]
        for item in self.codes:
            print(item)
            df=pd.read_csv('./data/'+item+'.csv')
            df['return']=df['收盘'].pct_change()
            #累计涨跌幅的序列
            initial_price = df['收盘'].iloc[0]
            df['return_sum']=(df['收盘'] - initial_price) / initial_price
            #create a new dataframe,and copy df['return'],df['close'] to it
            df_new=pd.DataFrame()

            df_new['close']=df['收盘']
            df_new['return'] = df['return']
            df_new['return_sum'] = df['return_sum']
            df_new['date'] = df['日期']
            df_new.to_csv('./data_return/'+item+'_return.csv')
            res=self.analysis(df,item)
            res_analysis.append(res)
        print(res_analysis)
        self.plot(res_analysis)
    def analysis(self,df,symbol):
        # 计算最终收益率
        final_return = df['return_sum'].iloc[-1]
        print(f"最终收益率: {final_return:.2%}")

        # 计算最大回撤
        df['peak'] = df['return_sum'].cummax()
        df['drawdown'] = df['peak'] - df['return_sum']
        max_drawdown = df['drawdown'].max()
        print(f"最大回撤: {max_drawdown:.2%}")

        # 计算波动率
        volatility = df['return'].std() * np.sqrt(252)  # 假设每年有252个交易日
        print(f"波动率: {volatility:.2%}")

        return [symbol,final_return,max_drawdown,volatility]
    def plot(self,data):
        # 转换为 DataFrame
        columns = ['Stock Code', 'Final Return', 'Max Drawdown', 'Volatility']
        df = pd.DataFrame(data, columns=columns)

        # 设置股票代码为索引
        df.set_index('Stock Code', inplace=True)

        # 绘制柱状图
        fig, ax = plt.subplots(1, 3, figsize=(15, 5))

        # Final Return
        df['Final Return'].plot(kind='bar', ax=ax[0], color='skyblue')
        ax[0].set_title('Final Return')
        ax[0].set_xlabel('Stock Code')
        ax[0].set_ylabel('Return')

        # Max Drawdown
        df['Max Drawdown'].plot(kind='bar', ax=ax[1], color='orange')
        ax[1].set_title('Max Drawdown')
        ax[1].set_xlabel('Stock Code')
        ax[1].set_ylabel('Drawdown')

        # Volatility
        df['Volatility'].plot(kind='bar', ax=ax[2], color='green')
        ax[2].set_title('Volatility')
        ax[2].set_xlabel('Stock Code')
        ax[2].set_ylabel('Volatility')

        plt.tight_layout()
        plt.show()
if __name__ == '__main__':
    ra=ReturnAnalysis()
    ra.get_return()


# 综合分析
# 股票代码 002594：最终收益率较高，最大回撤和波动率处于中等水平，表明该股票在考察期间表现较好，收益较高，但波动风险也相对较大。
# 股票代码 000905：最终收益率较低，但最大回撤和波动率较小，表明该股票在考察期间表现较为稳健，虽然收益不高，但风险较低。
# 股票代码 002607：最终收益率最低，最大回撤和波动率最大，表明该股票在考察期间表现较差，不仅亏损较大，而且价格波动剧烈，风险较高。
